#!/usr/env/bin python
# -*- coding: utf-8 -*-

# @Time    : 2020-11-24 9:37 上午
# @Author  : yangdy
# @File    : 2020-11-24.222.完全二叉树的节点个数
# @Software: PyCharmCE
from tools.my_methods import *


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


class Solution:
    def countNodes(self, root: TreeNode) -> int:
        if not root:
            return 0
        dep = self.depth(root) - 1
        ll, rr = 1 << dep, (1 << (dep + 1)) - 1
        while ll < rr:
            mid = ll + (rr - ll + 1) // 2
            if self.exists(root, dep, mid):
                ll = mid
            else:
                rr = mid - 1
        return ll

    def exists(self, rt, deep, k):
        e = 1 << (deep - 1)
        while rt and e:
            if e & k > 0:
                rt = rt.right
            else:
                rt = rt.left
            e >>= 1
        return rt is not None

    # def countNodes(self, root: TreeNode) -> int:
    #     display(root)
    #     if not root:
    #         return 0
    #     ll = self.depth(root.left)
    #     rr = self.depth(root.right)
    #     if ll > rr:
    #         return self.countNodes(root.left) + (1 << rr)
    #     else:
    #         return self.countNodes(root.right) + (1 << ll)
    #
    def depth(self, root: TreeNode) -> int:
        depth = 0
        while root:
            depth += 1
            root = root.left
        return depth

    # def countNodes(self, root: TreeNode) -> int:
    #     if not root:
    #         return 0
    #     ll = self.countNodes(root.left)
    #     rr = self.countNodes(root.right)
    #     return ll + rr + 1


questions = [
    ([1, 2, 3, 4, 5, 6],),
    ([1, 2, 3, 4],),
]

work(questions, lambda q: Solution().countNodes(mk_tree(q[0], TreeNode)))
